In _try_flashinfer_sample_without_penalty, logits.div_(temperature) is performed in-place. If the subsequent flashinfer sampling function returns None (e.g., because flashinfer is not available or fails), the function returns None, and the caller falls back to standard sampling. However, logits has already been modified in-place, so the caller will divide logits by the temperature again, leading to incorrect sampling probabilities. We should restore logits by multiplying back by temperature if the flashinfer call returns None.

The _add_rms_norm_fwd_fused kernel uses a loop over N with BLOCK_SIZE steps. However, N > BLOCK_SIZE is explicitly forbidden and raises a RuntimeError in the Python wrapper. Since N <= BLOCK_SIZE is guaranteed, the loop is completely redundant. More importantly, storing X and then reloading it in the second loop causes unnecessary global memory read/write operations, which is extremely slow for memory-bandwidth-bound kernels like RMSNorm. We can load X only once and avoid the loop entirely.

@triton.jit
def _add_rms_norm_fwd_fused(
    X,
    R,
    Y,
    W,
    x_stride0,
    x_stride1,
    r_stride0,
    r_stride1,
    y_stride0,
    y_stride1,
    N,
    eps,
    HAS_WEIGHT: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    row = tl.program_id(0)
    X += row * x_stride0
    R += row * r_stride0
    Y += row * y_stride0

    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < N
    x = tl.load(X + cols * x_stride1, mask=mask, other=0.0).to(tl.float32)
    r = tl.load(R + cols * r_stride1, mask=mask, other=0.0).to(tl.float32)
    x = x + r
    tl.store(X + cols * x_stride1, x.to(X.dtype.element_ty), mask=mask)

    var = tl.sum(x * x, axis=0) / N
    rstd = 1 / tl.sqrt(var + eps)
    y = x * rstd
    if HAS_WEIGHT:
        w = tl.load(W + cols, mask=mask).to(tl.float32)
        y *= w
    tl.store(Y + cols * y_stride1, y.to(Y.dtype.element_ty), mask=mask)

The in-place division logits.div_(b_temperatures.view((-1, 1))) is performed before checking if the request is greedy and doesn't need logprobs. For greedy sampling, temperature scaling does not affect the argmax result. Performing this division on a large logits tensor (e.g., vocab size 150k+) is a significant waste of GPU compute and memory bandwidth. Reordering the greedy check to return early before performing the division will improve performance.

In moe_align_fused_kernel, if ZERO_EXPERT_TOKEN_NUM is True, every program block in the grid will concurrently write 0 to expert_token_num_ptr. While functionally correct, this is redundant and causes unnecessary memory write contention. Adding and token_block == 0 to the condition ensures only the first block performs the zeroing.

In pack_gdn_decode_inputs, z_raw is loaded in the Triton kernel using flattened indexing (z_raw + row * stride_z_b + qkv_offsets), which assumes that the last two dimensions of z_raw are contiguous. However, there is no assertion in the Python wrapper to ensure z_raw is contiguous. If a non-contiguous tensor is passed, it will silently produce incorrect results. We should add an assertion to prevent this.

In conv_pack_gdn_decode_inputs, z_raw is loaded in the Triton kernel using flattened indexing, which assumes that the last two dimensions of z_raw are contiguous. However, there is no assertion in the Python wrapper to ensure z_raw is contiguous. If a non-contiguous tensor is passed, it will silently produce incorrect results. We should add an assertion to prevent this.

The manual implementation of sigmoid 1.0 / (1.0 + tl.exp(-gate)) is used. Triton has a built-in tl.sigmoid function which is cleaner, more readable, and can be better optimized by the compiler.

The manual implementation of sigmoid 1.0 / (1.0 + tl.exp(-gate_vals)) is used. Triton has a built-in tl.sigmoid function which is cleaner, more readable, and can be better optimized by the compiler.

The manual implementation of sigmoid 1.0 / (1.0 + tl.exp(-gate_vals)) is used. Triton has a built-in tl.sigmoid function which is cleaner, more readable, and can be better optimized by the compiler.